import random
import pandas as pd
import gradio as gr
import plotly.graph_objects as go

import warnings


warnings.filterwarnings("ignore")




pth = r".\data\2020中国投入产出表.xlsx"
data = pd.read_excel(pth, index_col=0, sheet_name=None)


def random_rgba_color():
    # 随机生成 RGB 分量，范围是 0 到 255
    red = random.randint(0, 255)
    green = random.randint(0, 255)
    blue = random.randint(0, 255)
    # 随机生成透明度分量，范围是 0 到 1 之间，保留一位小数
    alpha = round(random.uniform(0, 1), 1)
    # 构建并返回 rgba 颜色字符串
    return f'rgba({red}, {green}, {blue}, {alpha})'


def random_hex_color():
    """
    此函数用于随机生成一个十六进制颜色代码。
    返回:
        str: 一个以 '#' 开头的 6 位十六进制颜色代码。
    """
    # 生成 6 位十六进制数字
    hex_chars = '0123456789ABCDEF'
    color = ''.join(random.choice(hex_chars) for _ in range(6))
    # 拼接成完整的十六进制颜色代码
    return f'#{color}'


def show_complex_io(data_df, filter_ratio, select_items, select_cols, img_name):
    fill_df = data_df.copy()
    fill_df = fill_df.loc[select_items, select_cols]
    fill_df.columns = [f"O_{col}" for col in fill_df.columns]
    v = fill_df.sum().sum() * filter_ratio
    print('filter_ratio:', filter_ratio, fill_df.sum().sum(), v)
    source_li =[]
    target_li = []
    value_li = []
    source_color_li = []
    target_color_li = []
    link_color_li = []

    select_col_items = [f"O_{col}" for col in select_items]
    
    for col in fill_df.columns:
        col_df = fill_df[col]
        col_df = col_df[col_df > v]
        for index in col_df.index:
            if index in select_items or col in select_col_items:
                source_li.append(index)
                target_li.append(col)
                value_li.append(col_df[index])
                source_color_li.append(random_hex_color())
                target_color_li.append(random_hex_color())
                link_color_li.append(random_rgba_color())
        
    # 示例数据
    data = {
        'source': source_li,
        'target': target_li,
        'value': value_li,
        'source_color': source_color_li,
        'target_color': target_color_li,
        'link_color': link_color_li
    }
    df = pd.DataFrame(data)
    # 生成节点列表
    unique_nodes = list(set(df['source'].tolist() + df['target'].tolist()))
    node_indices = {node: index for index, node in enumerate(unique_nodes)}
    # 转换源和目标节点为索引
    df['source_index'] = df['source'].map(node_indices)
    df['target_index'] = df['target'].map(node_indices)
    # 创建 Sankey 图对象
    fig = go.Figure(data=[go.Sankey(
        node=dict(
            pad=15,
            thickness=20,
            line=dict(color="black", width=0.5),
            label=unique_nodes,
            color=[random_hex_color() for _ in range(len(unique_nodes))]
        ),
        link=dict(
            source=df['source_index'],
            target=df['target_index'],
            value=df['value'],
            color=df['link_color']
        )
    )])

    # 设置图表布局
    fig.update_layout(
        title_text=img_name,
        font_size=12,
        paper_bgcolor='rgba(0,0,0,0)',  # 透明背景
        plot_bgcolor='rgba(0,0,0,0)',
        height=900,
    )
    # fig.write_html('sankey.html')
    # 显示图表
    return fig, f"规则：显示流量大于{v:.2f}亿元"
    

class ShowData:
    def __init__(self):
        self.data = data
        self.matrix = data['中间消耗矩阵'] / 10000
        self.output = data['支出法矩阵'] / 10000
        self.input = data['收入法矩阵'] / 10000
        self.data_set = {'中间消耗': self.matrix, '最终需求': self.output, '最初投入': self.input}
        self.unit = '亿元'
        self.items = self.matrix.columns.to_list()
        self.columns = self.matrix.columns.to_list()
    
    def change_matrix(self, matrix):
        """
        此函数用于根据用户选择的矩阵名称，更新行节点和列节点的可选列表。
        它会从数据集中获取所选矩阵，并将其索引和列名作为新的可选节点。

        参数:
            matrix (str): 用户选择的矩阵名称，如 '中间消耗', '最终需求', '最初投入'。

        返回:
            tuple: 包含两个 Gradio 的 CheckboxGroup 对象，分别用于行节点和列节点的选择。
        """
        # 从数据集中获取用户选择的矩阵
        select_matrix = self.data_set[matrix]
        # 打印所选矩阵，用于调试
        print('用户选择矩阵 > select_matrix:', select_matrix)
        # 更新行节点可选列表
        self.items = select_matrix.index.to_list()
        # 更新列节点可选列表
        self.columns = select_matrix.columns.to_list()
        # 返回两个复选框组，分别用于行节点和列节点的选择，初始值为全选
        return gr.CheckboxGroup(
            # 行节点复选框组，选项为所选矩阵的索引
            choices=select_matrix.index.to_list(), value=select_matrix.index.to_list()), gr.CheckboxGroup(
                # 列节点复选框组，选项为所选矩阵的列名
                choices=select_matrix.columns.to_list(), value=select_matrix.columns.to_list())
            
    def plot_complex_io(self, filter_ratio, items, select_matrix, col_items):
        """
        此函数用于根据用户选择的过滤比例、行节点、矩阵和列节点，绘制复杂的投入产出桑基图。
        参数:
            filter_ratio (float): 过滤比例，用于筛选流量数据。
            items (list): 选择的行节点列表。
            select_matrix (str): 用户选择的矩阵名称，如 '中间消耗', '最终需求', '最初投入'。
            col_items (list): 选择的列节点列表。
        返回:
            tuple: 包含绘制的桑基图对象和显示过滤规则的文本信息。
        """
        # 从数据集中获取用户选择的矩阵
        matrix = self.data_set[select_matrix]
        # 调用 show_complex_io 函数，传入所选矩阵、过滤比例、行节点和列节点，绘制桑基图
        return show_complex_io(
            matrix, 
            filter_ratio, 
            items, 
            col_items, 
            img_name=select_matrix
            )
    
    def item_select_all_fn(self):
        """
        此函数用于全选行节点的复选框组。
        它会返回一个 Gradio 的 CheckboxGroup 对象，其中所有选项都被选中。

        返回:
            gr.CheckboxGroup: 一个复选框组对象，其所有选项都被选中。
        """
        return gr.CheckboxGroup(choices=self.items, value=self.items)

    def item_deselect_all_fn(self):
        """
        此函数用于全不选行节点的复选框组。
        它会返回一个 Gradio 的 CheckboxGroup 对象，其中所有选项都未被选中。

        返回:
            gr.CheckboxGroup: 一个复选框组对象，其所有选项都未被选中。
        """
        return gr.CheckboxGroup(choices=self.items, value=[])
    
    def column_item_select_all_fn(self):
        """
        此函数用于全选列节点的复选框组。
        它会返回一个 Gradio 的 CheckboxGroup 对象，其中所有选项都被选中。
      
        返回:
            gr.CheckboxGroup: 一个复选框组对象，其所有选项都被选中。
        """
        # 返回一个复选框组，选项为列节点列表，且所有选项都被选中
        return gr.CheckboxGroup(choices=self.columns, value=self.columns)

    def column_item_deselect_all_fn(self):
        """
        此函数用于全不选列节点的复选框组。
        它会返回一个 Gradio 的 CheckboxGroup 对象，其中所有选项都未被选中。

        返回:
            gr.CheckboxGroup: 一个复选框组对象，其所有选项都未被选中。
        """
        # 返回一个复选框组，选项为列节点列表，所有选项都不被选中
        return gr.CheckboxGroup(choices=self.columns, value=[])
    
    def html(self):
        with gr.Blocks() as demo:
            gr.Markdown("# 投入产出表:流量分析")
            with gr.Row():
                matrix_select = gr.Radio(['中间消耗', '最终需求', '最初投入'], value='中间消耗', label="选择矩阵", scale=2)
                ratio = gr.Slider(0.0001, 0.01, value=0.005, step=0.0001, label="阈值忽略比例（占矩阵之和）", scale=2)
                text = gr.Textbox(value='', scale=1)
            with gr.Row():
                if_ratio = gr.Checkbox(label="计算比例", value=False, scale=1)
                plot_bt = gr.Button("生成流量图", scale=2)
            with gr.Tab('流量图'):
                complex_plot = gr.Plot(label="复杂网络")
            
            with gr.Tab("行节点"):
                items = gr.CheckboxGroup(choices=self.items, label="选择节点", value=self.items, interactive=True)
                with gr.Row():
                    items_select_all = gr.Button("全选")
                    items_deselect_all = gr.Button("全不选")
                    items_select_all.click(fn=self.item_select_all_fn, inputs=[], outputs=items)
                    items_deselect_all.click(fn=self.item_deselect_all_fn, inputs=[], outputs=items)
            
            with gr.Tab("列节点"):
                col_items = gr.CheckboxGroup(choices=self.items, label="选择节点", value=self.items, interactive=True)
                with gr.Row():
                    col_items_select_all = gr.Button("全选")
                    col_items_deselect_all = gr.Button("全不选")
                    col_items_select_all.click(fn=self.column_item_select_all_fn, inputs=[], outputs=col_items)
                    col_items_deselect_all.click(fn=self.column_item_deselect_all_fn, inputs=[], outputs=col_items)

            matrix_select.change(fn=self.change_matrix, inputs=matrix_select, outputs=[items, col_items])
            
            gr.on(
                triggers=[ratio.change, plot_bt.click], 
                outputs=[complex_plot, text], 
                fn=self.plot_complex_io, 
                inputs=[ratio, items, matrix_select, col_items]
                )
        return demo


sd = ShowData()
sd.html().launch()